import sys
import os

import cv2
from PyQt5 import QtCore, QtGui, QtWidgets
from PyQt5.QtCore import QObject, pyqtSignal, QThread, QStringListModel, Qt, QItemSelectionModel, QEvent, QRegExp
from PyQt5.QtGui import QImage, QPixmap, QBrush, QColor, QRegExpValidator
from PyQt5.QtWidgets import QMainWindow, QPlainTextEdit, QFileSystemModel, QListView, QMenu, QMessageBox, \
    QStyledItemDelegate

from ui.ui_main import Ui_MainWindow
from logconfig.logconfig import logger


def get_video_frame_count(video_path):
    """
    获取视频的总帧数
    :param video_path: 视频文件路径
    :return: 视频的总帧数
    """
    # 打开视频文件
    cap = cv2.VideoCapture(video_path)

    if not cap.isOpened():
        logger.error(f"无法打开视频文件: {video_path}")
        return -1

    # 获取总帧数
    frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))

    # 释放视频对象
    cap.release()

    return frame_count


def is_video_file(file_path):
    """判断文件是否为视频文件"""
    video_extensions = [".mp4", ".avi", ".mkv", ".mov", ".flv", ".wmv"]  # 支持的视频扩展名
    return os.path.isfile(file_path) and any(file_path.lower().endswith(ext) for ext in video_extensions)


def is_text_file(file_path):
    """判断文件是否为视频文件"""
    text_extensions = [".txt"]  # 支持的视频扩展名
    return os.path.isfile(file_path) and any(file_path.lower().endswith(ext) for ext in text_extensions)


def get_frame_pic(video_path, frame_count, frame_number):
    """显示指定帧"""
    # 检查帧号是否在有效范围内
    if frame_number < 1 or frame_number > frame_count:
        logger.error(f"帧号 {frame_number} 超出有效范围（1 - {frame_count}）")
        return None

    # 打开视频文件
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        logger.error(f"无法打开视频文件: {video_path}")
        return None
    # 跳转到指定帧
    cap.set(cv2.CAP_PROP_POS_FRAMES, frame_number - 1)  # 帧号从 0 开始
    ret, frame = cap.read()

    if not ret:
        logger.error(f"无法读取帧: {frame_number}")
        cap.release()
        return None
    cap.release()
    # 将 OpenCV 的 BGR 格式转换为 RGB 格式
    frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

    # 将 OpenCV 图片转换为 PyQt5 图片
    height, width, channel = frame_rgb.shape
    bytes_per_line = 3 * width
    q_image = QImage(frame_rgb.data, width, height, bytes_per_line, QImage.Format_RGB888)

    # 将 QImage 转换为 QPixmap 并显示
    pixmap = QPixmap.fromImage(q_image)
    return pixmap


class Stream(QObject):
    newText = pyqtSignal(str)

    def write(self, text):
        self.newText.emit(str(text))


class CustomStringListModel(QStringListModel):
    def __init__(self, parent=None):
        super().__init__(parent)
        self._background_colors = {}  # 存储背景颜色

    def setData(self, index, value, role=Qt.EditRole):
        if role == Qt.BackgroundRole:
            self._background_colors[index] = value
            self.dataChanged.emit(index, index, [role])
            return True
        return super().setData(index, value, role)

    def data(self, index, role=Qt.DisplayRole):
        if role == Qt.BackgroundRole:
            return self._background_colors.get(index, None)
        return super().data(index, role)

    def clearBackgroundColors(self):
        """清除所有背景颜色"""
        self._background_colors.clear()
        self.dataChanged.emit(self.index(0), self.index(self.rowCount() - 1), [Qt.BackgroundRole])


class GetFramePicThread(QThread):
    frame_pic_ready = pyqtSignal(object)

    def __init__(self, video_path, frame_count, frame_number):
        super().__init__()
        self.video_path = video_path
        self.frame_count = frame_count
        self.frame_number = frame_number
        self._stop_flag = False  # 添加停止标志

    def run(self):
        try:
            while not self._stop_flag:  # 检查停止标志
                img = get_frame_pic(self.video_path, self.frame_count, self.frame_number)
                if img is not None:
                    self.frame_pic_ready.emit(img)
                else:
                    logger.error("未能获取到指定帧的图片")
                break  # 处理完一帧后退出循环
        except Exception as e:
            logger.error(f"获取帧图片时发生错误: {e}")

    def stop(self):
        self._stop_flag = True  # 设置停止标志
        self.wait()  # 等待线程结束


class MainWindow(QMainWindow, Ui_MainWindow):
    def __init__(self):
        super().__init__()
        self.frame_number = None
        self.thread = None  # 初始化线程为 None
        self.setupUi(self)
        self.setWindowTitle("性能测试数帧工具")
        self.setFixedSize(1307, 959)
        self.init_log_editor()
        self.cal_hz = 0
        self.folder_path = ''
        self.video_path = ''
        self.frame_count = 0
        self.model = QFileSystemModel()
        self.listmodel = CustomStringListModel()
        self.start_frame = None
        self.end_frame = None
        self.time_res = None
        self.init_ui()
        # 为 frame_list_view 安装事件过滤器
        self.frame_list_view.installEventFilter(self)

    def init_ui(self):
        """初始化UI界面"""
        self.init_radio_btn()
        self.init_frame_pic()
        self.frame_editor.setDisabled(True)
        self.jump_frame_btn.setDisabled(True)
        self.cal_time_btn.setDisabled(True)
        self.frame_res.setReadOnly(True)

    def init_frame_pic(self):
        """初始化图片显示区域"""
        self.frame_pic.setAlignment(Qt.AlignmentFlag.AlignCenter)  # 图片居中显示
        self.frame_pic.setStyleSheet("background-color: white;")  # 设置背景色（可选）
        self.frame_pic_layout.addWidget(self.frame_pic, 0, 0)  # 添加到布局中
        self.frame_pic_layout.setRowStretch(0, 1)  # 设置拉伸因子
        self.frame_pic_layout.setColumnStretch(0, 1)

    def init_radio_btn(self):
        """初始化帧数计算选择按钮"""
        self.choose_60hz_btn.setChecked(False)
        self.choose_120hz_btn.setChecked(True)
        self.cal_hz = 120
        self.choose_240hz_btn.setChecked(False)

    def init_log_editor(self):
        """初始化日志编辑器"""
        self.log_editor.setReadOnly(True)
        self.log_editor.setLineWrapMode(QPlainTextEdit.WidgetWidth)
        logger.add(self.write_log)
        sys.stdout = Stream(newText=self.on_update_text)
        sys.stderr = Stream(newText=self.on_update_text)

    def init_file_tree_view(self, folder_path):
        """初始化文件树视图"""
        try:
            if not folder_path:
                logger.error("文件夹路径为空")
                return
            self.model.setRootPath(folder_path)
            self.file_tree_view.setModel(self.model)
            self.file_tree_view.setRootIndex(self.model.index(folder_path))
            self.file_tree_view.setHeaderHidden(True)
            for i in range(1, self.model.columnCount()):
                self.file_tree_view.hideColumn(i)
        except Exception as e:
            logger.error(f"初始化文件树视图时发生错误: {e}")

    def init_frame_list_view(self, video_path, frame_count):
        """初始化帧列表视图"""
        try:
            self.frame_list_view.setEditTriggers(QListView.NoEditTriggers)
            data_list = [f"第{i + 1}帧" for i in range(frame_count)]
            self.listmodel.setStringList(data_list)
            self.frame_list_view.setModel(self.listmodel)
            self.load_markers(video_path)
            logger.info("起始帧:{},结束帧:{}".format(self.start_frame, self.end_frame))
            self.listmodel.clearBackgroundColors()
            # 设置起始帧和结束帧的颜色
            if self.start_frame is not None:
                index = self.listmodel.index(self.start_frame - 1, 0)
                self.listmodel.setData(index, QColor('green'), Qt.BackgroundRole)
            if self.end_frame is not None:
                index = self.listmodel.index(self.end_frame - 1, 0)
                self.listmodel.setData(index, QColor('red'), Qt.BackgroundRole)

            if self.time_res is not None:
                self.frame_res.clear()
                self.frame_res.setText(str(self.time_res))
                self.frame_res.update()

            if self.listmodel.rowCount() > 0:
                first_index = self.listmodel.index(0)
                self.frame_list_view.setCurrentIndex(first_index)
                self.frame_list_view.selectionModel().select(first_index, QItemSelectionModel.SelectionFlag.Select)
            self.frame_list_view.selectionModel().selectionChanged.connect(self.on_selection_changed)
        except Exception as e:
            logger.error(f"初始化帧列表视图时发生错误: {e}")

    def update_pic(self, frame_number):
        """更新图片"""
        logger.info(f"视频路径: {self.video_path}")
        logger.info(f"总帧数: {self.frame_count}")
        logger.info(f"选中的帧号: {frame_number}")

        if not self.video_path or self.frame_count <= 0:
            logger.error("视频路径或总帧数无效")
            return

        if self.thread and self.thread.isRunning():
            self.thread.stop()  # 停止上一个线程

        self.thread = GetFramePicThread(self.video_path, self.frame_count, frame_number)
        self.thread.frame_pic_ready.connect(self.show_frame_pic)
        self.thread.start()

    def show_frame_pic(self, img):
        """显示帧图片"""
        try:
            if img.isNull():
                logger.error("未能获取到指定帧的图片")
                return

            # 获取 frame_pic 的当前大小
            label_size = self.frame_pic.size()

            # 等比例缩放图片
            scaled_pixmap = img.scaled(
                label_size,
                Qt.AspectRatioMode.KeepAspectRatio,  # 等比例缩放
                Qt.TransformationMode.SmoothTransformation  # 平滑缩放
            )

            # 设置图片到 QLabel
            self.frame_pic.setPixmap(scaled_pixmap)
            self.frame_pic.update()  # 强制刷新
        except Exception as e:
            logger.error(f"显示帧图片时发生错误: {e}")

    def on_selection_changed(self, selected, deselected):
        """当选中项变化时触发"""
        try:
            selected_indexes = selected.indexes()
            if selected_indexes:
                selected_item = selected_indexes[0].data()
                selected_item_index = selected_indexes[0].row()
                logger.info(f"选中的项: {selected_item}")
                logger.info(f"选中项的索引: {selected_item_index}")
                self.update_pic(selected_item_index + 1)  # 帧号从 1 开始
            else:
                logger.info("未选中任何项")
        except Exception as e:
            logger.error(f"报错信息如下: {str(e)}")

    def on_update_text(self, text):
        """更新日志"""
        cursor = self.log_editor.textCursor()
        cursor.movePosition(cursor.End)
        cursor.insertText(text)
        self.log_editor.setTextCursor(cursor)
        self.log_editor.ensureCursorVisible()

    def write_log(self, message):
        """写入日志"""
        self.log_editor.appendPlainText(str(message).strip())

    def click60hz_btn(self):
        """选中60hz计算"""
        self.cal_hz = 60
        logger.info(self.cal_hz)

    def click120hz_btn(self):
        """选中120hz计算"""
        self.cal_hz = 120
        logger.info(self.cal_hz)

    def click240hz_btn(self):
        """选中240hz计算"""
        self.cal_hz = 240
        logger.info(self.cal_hz)

    def choose_folder(self):
        """选择文件夹"""
        try:
            from PyQt5.QtWidgets import QFileDialog
            folder_path = QFileDialog.getExistingDirectory(self, '选择文件夹')
            if folder_path and os.path.isdir(folder_path):
                logger.info(f'所选文件夹路径: {folder_path}')
                self.folder_path = folder_path
                self.init_file_tree_view(self.folder_path)
            else:
                logger.error('无效的文件夹路径')
        except Exception as e:
            logger.error(f"选择文件夹时发生错误: {e}")

    def choose_file(self, index):
        """选择文件"""
        try:
            if self.thread and self.thread.isRunning():
                self.thread.stop()  # 停止上一个线程

            self.start_frame = None
            self.end_frame = None
            self.time_res = None
            self.frame_res.clear()
            self.frame_editor.clear()
            self.frame_editor.setPlaceholderText("请输入跳转帧号")
            # 设置正则表达式验证器，允许空字符串或正整数
            reg_exp = QRegExp("^[1-9]\\d*|$")  # 正则表达式：允许空字符串或以 1-9 开头的数字
            validator = QRegExpValidator(reg_exp, self)
            self.frame_editor.setValidator(validator)
            file_choose_path = self.model.filePath(index)
            logger.info(file_choose_path)
            if is_video_file(file_choose_path):
                self.frame_editor.setDisabled(False)
                self.jump_frame_btn.setDisabled(False)
                self.cal_time_btn.setDisabled(False)
                self.video_path = file_choose_path
                frame_count = get_video_frame_count(file_choose_path)
                self.frame_count = frame_count
                logger.info("视频总帧数为:{}".format(frame_count))
                self.init_frame_list_view(file_choose_path, frame_count)
            else:
                self.frame_editor.setDisabled(True)
                self.jump_frame_btn.setDisabled(True)
                self.cal_time_btn.setDisabled(True)
                if is_text_file(file_choose_path):
                    self.listmodel.removeRows(0, self.listmodel.rowCount())
                    if self.thread and self.thread.isRunning():
                        self.thread.stop()  # 停
                    self.frame_pic.clear()
                    self.frame_res.clear()
                    self.load_txt(file_choose_path)

        except Exception as e:
            logger.error(f"选择文件时发生错误: {e}")

    def eventFilter(self, obj, event):
        """事件过滤器，处理右键点击事件"""
        if obj == self.frame_list_view and event.type() == QEvent.ContextMenu:
            # 创建右键菜单
            menu = QMenu(self)
            # 添加菜单项
            action1 = menu.addAction("标记起始帧")
            action2 = menu.addAction("标记结束帧")
            action3 = menu.addAction("取消标记帧")
            action4 = menu.addAction("跳转到上一帧")
            action5 = menu.addAction("跳转到下一帧")

            # 连接菜单项的触发信号到相应的槽函数
            action1.triggered.connect(self.action1_triggered)
            action2.triggered.connect(self.action2_triggered)
            action3.triggered.connect(self.action3_triggered)
            action4.triggered.connect(self.action4_triggered)
            action5.triggered.connect(self.action5_triggered)

            # 显示右键菜单
            menu.exec_(event.globalPos())
            return True
        return super().eventFilter(obj, event)

    def action1_triggered(self):
        """标记起始帧"""
        try:
            selected_indexes = self.frame_list_view.selectedIndexes()
            if selected_indexes:
                frame_number = selected_indexes[0].row()
                if self.end_frame is not None and (
                        frame_number + 1 > self.end_frame or frame_number + 1 == self.end_frame):
                    QMessageBox.warning(self, "错误", "起始帧不能在结束帧之后！")
                    return
                self.start_frame = frame_number + 1
                self.update_marker_colors()
                self.save_markers()
        except Exception as e:
            logger.error(str(e))

    def action2_triggered(self):
        """标记结束帧"""
        try:
            selected_indexes = self.frame_list_view.selectedIndexes()
            if selected_indexes:
                frame_number = selected_indexes[0].row()
                if self.start_frame is not None and (
                        frame_number + 1 < self.start_frame or frame_number + 1 == self.start_frame):
                    QMessageBox.warning(self, "错误", "结束帧不能在起始帧之前！")
                    return
                self.end_frame = frame_number + 1
                self.update_marker_colors()
                self.save_markers()
        except Exception as e:
            logger.error(str(e))

    def action3_triggered(self):
        """取消标记帧"""
        try:
            current_index = self.frame_list_view.currentIndex()
            if self.start_frame is not None and self.end_frame is not None:
                start_frame_index = self.listmodel.index(self.start_frame - 1)
                end_frame_index = self.listmodel.index(self.end_frame - 1)
                if current_index == start_frame_index:
                    self.start_frame = None
                    self.listmodel.setData(current_index, QColor('white'), Qt.BackgroundRole)
                    self.frame_res.clear()
                    self.time_res = None
                    self.save_markers()
                elif current_index == end_frame_index:
                    self.end_frame = None
                    self.listmodel.setData(current_index, QColor('white'), Qt.BackgroundRole)
                    self.frame_res.clear()
                    self.time_res = None
                    self.save_markers()
                else:
                    pass
            elif self.start_frame is not None and self.end_frame is None:
                start_frame_index = self.listmodel.index(self.start_frame - 1)
                if current_index == start_frame_index:
                    self.start_frame = None
                    self.listmodel.setData(current_index, QColor('white'), Qt.BackgroundRole)
                    self.frame_res.clear()
                    self.time_res = None
                    self.save_markers()
                else:
                    pass
            elif self.start_frame is None and self.end_frame is not None:
                end_frame_index = self.listmodel.index(self.end_frame - 1)
                if current_index == end_frame_index:
                    self.end_frame = None
                    self.listmodel.setData(current_index, QColor('white'), Qt.BackgroundRole)
                    self.frame_res.clear()
                    self.time_res = None
                    self.save_markers()
                else:
                    pass
        except Exception as e:
            logger.error(str(e))

    def action5_triggered(self):
        """跳转到下一帧"""
        try:
            current_index = self.frame_list_view.currentIndex()
            if self.start_frame is not None and self.end_frame is not None:
                start_frame_index = self.listmodel.index(self.start_frame - 1)
                end_frame_index = self.listmodel.index(self.end_frame - 1)
                next_index = current_index
                if current_index < start_frame_index:
                    next_index = start_frame_index
                elif start_frame_index < current_index < end_frame_index:
                    next_index = end_frame_index

                if next_index.isValid():
                    self.frame_list_view.setCurrentIndex(next_index)
                    self.frame_list_view.selectionModel().select(next_index, QItemSelectionModel.SelectionFlag.Select)
                logger.info("跳转到下一帧")
            elif self.start_frame is not None and self.end_frame is None:
                start_frame_index = self.listmodel.index(self.start_frame - 1)
                next_index = current_index
                if current_index < start_frame_index:
                    next_index = start_frame_index
                if next_index.isValid():
                    self.frame_list_view.setCurrentIndex(next_index)
                    self.frame_list_view.selectionModel().select(next_index, QItemSelectionModel.SelectionFlag.Select)
                logger.info("跳转到下一帧")

            elif self.start_frame is None and self.end_frame is not None:
                end_frame_index = self.listmodel.index(self.end_frame - 1)
                next_index = current_index
                if current_index < end_frame_index:
                    next_index = end_frame_index
                if next_index.isValid():
                    self.frame_list_view.setCurrentIndex(next_index)
                    self.frame_list_view.selectionModel().select(next_index, QItemSelectionModel.SelectionFlag.Select)
                logger.info("跳转到下一帧")
        except Exception as e:
            logger.error(str(e))

    def action4_triggered(self):
        """跳转到上一帧"""
        try:
            current_index = self.frame_list_view.currentIndex()
            if self.start_frame is not None and self.end_frame is not None:
                start_frame_index = self.listmodel.index(self.start_frame - 1)
                end_frame_index = self.listmodel.index(self.end_frame - 1)
                prev_index = current_index
                if end_frame_index < prev_index:
                    prev_index = end_frame_index
                elif start_frame_index < prev_index < end_frame_index:
                    prev_index = start_frame_index

                if prev_index.isValid():
                    self.frame_list_view.setCurrentIndex(prev_index)
                    self.frame_list_view.selectionModel().select(prev_index, QItemSelectionModel.SelectionFlag.Select)
                logger.info("跳转到上一帧")
            elif self.start_frame is not None and self.end_frame is None:
                prev_index = current_index
                start_frame_index = self.listmodel.index(self.start_frame - 1)
                if start_frame_index < prev_index:
                    prev_index = start_frame_index
                    if prev_index.isValid():
                        self.frame_list_view.setCurrentIndex(prev_index)
                        self.frame_list_view.selectionModel().select(prev_index,
                                                                     QItemSelectionModel.SelectionFlag.Select)
                    logger.info("跳转到上一帧")
            elif self.start_frame is None and self.end_frame is not None:
                end_frame_index = self.listmodel.index(self.end_frame - 1)
                prev_index = current_index
                if end_frame_index < prev_index:
                    prev_index = end_frame_index
                    if prev_index.isValid():
                        self.frame_list_view.setCurrentIndex(prev_index)
                        self.frame_list_view.selectionModel().select(prev_index,
                                                                     QItemSelectionModel.SelectionFlag.Select)
                    logger.info("跳转到上一帧")

        except Exception as e:
            logger.error(str(e))

    def load_markers(self, video_path):
        """从txt文件中加载标记信息"""
        txt_path = os.path.splitext(video_path)[0] + '.txt'
        if os.path.exists(txt_path):
            with open(txt_path, 'r', encoding='utf-8') as f:
                lines = f.readlines()
                for line in lines:
                    if line.startswith('起始帧:'):
                        frame_str = line.split(':')[1].strip()
                        self.start_frame = int(frame_str) if frame_str != 'null' else None
                    elif line.startswith('结束帧:'):
                        frame_str = line.split(':')[1].strip()
                        self.end_frame = int(frame_str) if frame_str != 'null' else None
                    elif line.startswith('时间(ms):'):
                        time_res_str = line.split(':')[1].strip()
                        self.time_res = float(time_res_str) if time_res_str != 'null' else None
        logger.info("起始帧:{},结束帧:{},时间(ms):{}".format(self.start_frame, self.end_frame, self.time_res))

    def load_txt(self, txt_path):
        """从txt文件中加载标记信息"""
        if os.path.exists(txt_path):
            with open(txt_path, 'r', encoding='utf-8') as f:
                lines = f.readlines()
                res_str = ""
                for line in lines:
                    res_str = res_str + line + "\n"
                self.frame_pic.setText(res_str)

    def save_markers(self):
        """将标记信息保存到txt文件中"""
        try:
            video_path = self.video_path  # 假设 self.current_video_path 是当前视频的路径
            txt_path = os.path.splitext(video_path)[0] + '.txt'
            with open(txt_path, 'w', encoding='utf-8') as f:
                start_frame_str = str(self.start_frame) if self.start_frame is not None else 'null'
                end_frame_str = str(self.end_frame) if self.end_frame is not None else 'null'
                time_res_str = str(self.time_res) if self.time_res is not None else 'null'
                f.write(f'起始帧:{start_frame_str}\n')
                f.write(f'结束帧:{end_frame_str}\n')
                f.write(f'时间(ms):{time_res_str}\n')
        except Exception as e:
            logger.error(f"选择文件时发生错误: {e}")

    def update_marker_colors(self):
        """更新标记帧的颜色"""
        logger.info("起始帧:{},结束帧:{}".format(self.start_frame, self.end_frame))
        self.listmodel.clearBackgroundColors()
        # 设置起始帧和结束帧的颜色
        if self.start_frame is not None:
            index = self.listmodel.index(self.start_frame - 1)
            success = self.listmodel.setData(index, QColor('green'), Qt.BackgroundRole)
            if not success:
                logger.error("Failed to set data for start frame")
        if self.end_frame is not None:
            index = self.listmodel.index(self.end_frame - 1)
            success = self.listmodel.setData(index, QColor('red'), Qt.BackgroundRole)
            if not success:
                logger.error("Failed to set data for end frame")

    def cal_time(self):
        try:
            if self.start_frame is not None and self.end_frame is not None:
                if self.start_frame < self.end_frame:
                    time_res = ((self.end_frame - self.start_frame) / self.cal_hz) * 1000
                    time_res = round(time_res, 2)
                    self.time_res = time_res
                    logger.info(
                        "开始帧为:{},结束帧为:{},视频帧率为:{},计算时间为:{}ms".format(self.start_frame, self.end_frame,
                                                                                       self.cal_hz, self.time_res))
                    self.frame_res.clear()
                    self.frame_res.setText(str(self.time_res))
                    self.frame_res.update()
                    self.save_markers()
                else:
                    logger.error("非法计算")
        except Exception as e:
            logger.error(str(e))

    def get_frame_number(self):
        try:
            if self.frame_editor.text().isdigit():
                self.frame_number = int(self.frame_editor.text())
                logger.info("输入的查询帧号为:{}".format(self.frame_number))
        except Exception as e:
            logger.error(str(e))

    def jump_frame_pic(self):
        if self.frame_number is not None:
            if 1 <= self.frame_number <= self.frame_count:
                index = self.listmodel.index(self.frame_number - 1)
                if index.isValid():
                    self.frame_list_view.setCurrentIndex(index)
                    self.frame_list_view.selectionModel().select(index, QItemSelectionModel.SelectionFlag.Select)
            else:
                QMessageBox.warning(self, "错误", "搜索帧号超过范围！")
                return


if __name__ == '__main__':
    app = QtWidgets.QApplication(sys.argv)
    window = MainWindow()
    window.show()
    sys.exit(app.exec_())
